import os
import pymysql
from urllib.parse import urlparse
from langchain_core.tools import tool
from langchain_community.chat_models.tongyi import ChatTongyi
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser



def run_sql_query(query, db_path="root:123456@localhost:3306/sim_war_game"):
    """
    执行传入的SQL查询（适用于 MySQL）
    args:
    - query: 要执行的SQL
    - db_path: MySQL连接路径
    return：查询的结果（列表 of 字典）
    """
    # 解析 db_path，假设格式为 user:password@host:port/database
    try:
        # 支持格式：user:pass@host:port/dbname
        parsed = urlparse(f'mysql://{db_path}')
        netloc = parsed.netloc
        path = parsed.path.lstrip('/')

        # 解析用户名密码
        if '@' in netloc:
            auth, host_port = netloc.split('@')
            if ':' in auth:
                user, password = auth.split(':')
            else:
                user, password = auth, ''
        else:
            host_port = netloc
            user, password = 'root', ''  # 默认用户

        # 解析 host 和 port
        if ':' in host_port:
            host, port = host_port.split(':')
            port = int(port)
        else:
            host, port = host_port, 3306

        database = path or 'mysql'

    except Exception as e:
        raise ValueError(f"无法解析数据库路径: {db_path}") from e

    # 连接数据库并执行查询
    connection = None
    try:
        connection = pymysql.connect(
            host=host,
            port=port,
            user=user,
            password=password,
            database=database,
            charset='utf8mb4',
            cursorclass=pymysql.cursors.DictCursor  # 返回字典格式
        )
        with connection.cursor() as cursor:
            cursor.execute(query)
            result = cursor.fetchall()
        connection.commit()
        return result

    except Exception as e:
        raise RuntimeError(f"SQL执行失败: {e}") from e

    finally:
        if connection:
            connection.close()


def extract_schema_prompt(db_path="root:123456@localhost:3306/sim_war_game"):
    """
    从 MySQL 数据库路径中提取数据库模式信息。

    args:
    - db_path: MySQL 连接路径

    return:
    - 包含所有表、字段、约束及注释的格式化字符串
    """
    try:
        # 构造伪URL以便使用urlparse
        if not db_path.startswith("mysql://"):
            db_path = f"mysql://{db_path}"
        
        parsed = urlparse(db_path)
        
        username = parsed.username
        password = parsed.password
        host = parsed.hostname
        port = parsed.port or 3306
        database = parsed.path.lstrip('/')
        if not database:
            raise ValueError("缺少数据库名")

        # 建立数据库连接
        connection = pymysql.connect(
            host=host,
            port=port,
            user=username,
            password=password,
            database=database,
            charset='utf8mb4'
        )
        cursor = connection.cursor()

        prompt_lines = []

        # 获取所有表名及其注释
        cursor.execute("""
            SELECT TABLE_NAME, TABLE_COMMENT 
            FROM INFORMATION_SCHEMA.TABLES 
            WHERE TABLE_SCHEMA = %s AND TABLE_TYPE = 'BASE TABLE';
        """, (database,))
        tables = cursor.fetchall()

        for table_name, table_comment in tables:
            # 添加表名和表注释
            prompt_lines.append(f"Table: {table_name}")
            if table_comment and table_comment.strip():
                prompt_lines.append(f"  Comment: {table_comment.strip()}")

            # 获取列信息（含注释）
            cursor.execute("""
                SELECT 
                    COLUMN_NAME, 
                    DATA_TYPE, 
                    IS_NULLABLE, 
                    COLUMN_DEFAULT, 
                    COLUMN_KEY,
                    COLUMN_COMMENT
                FROM INFORMATION_SCHEMA.COLUMNS 
                WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s
                ORDER BY ORDINAL_POSITION;
            """, (database, table_name))
            
            columns = cursor.fetchall()
            for col in columns:
                name, col_type, is_nullable, dflt_value, column_key, col_comment = col
                desc = f" - {name}: {col_type}"
                if is_nullable == 'NO':
                    desc += " NOT NULL"
                if dflt_value is not None:
                    # 转义单引号
                    if isinstance(dflt_value, str):
                        dflt_value = dflt_value.replace("'", "''")
                    desc += f" DEFAULT '{dflt_value}'"
                if column_key == 'PRI':
                    desc += " PRIMARY KEY"
                if column_key == 'MUL':
                    desc += " INDEX"  # 简单标记外键相关索引
                prompt_lines.append(desc)

                # 添加字段注释（如果存在）
                if col_comment and col_comment.strip():
                    prompt_lines.append(f"   > {col_comment.strip()}")

            prompt_lines.append("")  # 表之间空行分隔

        cursor.close()
        connection.close()

        return "\n".join(prompt_lines)

    except Exception as e:
        raise RuntimeError(f"提取数据库模式失败: {e}") from e


def gen_system_prompt(prompt):
    """
    生成系统提示并调用语言模型生成SQL查询。

    args:
    - prompt: 用户输入的查询需求。

    return:
    - 生成的SQL查询语句。
    """

    # 定义提示模板
    prompt_template = ChatPromptTemplate.from_messages(
        [
            ("system", """
你是一个SQL专家，请根据提供的表结构信息和查询需求，来生成可执行的mysql支持的sql查询语句，你只能回复可执行的SQL内容，
生成的sql放在<SQL>标签内部

A：查询scene_id='a1b2c3d4-e5f6-7890-g1h2-i3j4k5l6m7n8'红方配置信息
Q：<SQL>select red_config from t_scene where scene_id="a1b2c3d4-e5f6-7890-g1h2-i3j4k5l6m7n8"</SQL>
"""),
            ("human", """
【数据库的结构】
{schema}

【查询需求】
{input}
"""),
        ]
    )

    # 用于获取数据库结构信息
    schema = extract_schema_prompt()

    # 使用提示模板生成提示词
    prompt_value = prompt_template.invoke({"input": prompt, "schema": schema})

    # 实例化一个大模型
    llm = ChatTongyi(
        model='qwen-plus',
        api_key="sk-da90821cf9174fbeb854011015c67aad"
    )

    # 调用语言模型生成SQL查询
    res = llm.invoke(prompt_value).content

    # 移除 <SQL> 标签
    sql_query = res.replace("<SQL>", "").replace("</SQL>", "")

    return sql_query


@tool
def load_red_config(user_prompt):
    """红方配置加载工具，将红方信息从MySQL数据库中查询出来

    args:
    - user_prompt: 用户输出的自然语言查询语句

    return: 执行SQL得到数据库中想要的己方配置信息
    """
    # 将用户的自然语言查询语句转为sql
    generated_sql = gen_system_prompt(user_prompt)

    # 执行SQL得到数据库中想要的己方配置信息
    response = run_sql_query(generated_sql)

    return response



if __name__ == "__main__":

    # 接入大模型
    model = ChatTongyi(
        model='qwen-plus',
        api_key="sk-da90821cf9174fbeb854011015c67aad"
    )

    # 创建一个提示词模板
    prompt = ChatPromptTemplate([
        ("system", "你是一个乐于助人的AI助手"),
        ("user", "查询scene_id='{scene_id}'红方配置信息")
    ])

    
    # 创建工具集
    tools = [load_red_config]

    # 把工具集注册到大模型
    model_with_tools = model.bind_tools(tools)

    # 构建一个链
    chain = prompt | model_with_tools

    fc =  chain.invoke({"scene_id": "b2a1d3c4-f5e6-7890-1234-567890abcdef"})

    for tool in fc.tool_calls:
        if tool.get('name') == "load_red_config":
            tool_msg = load_red_config.invoke(tool)
            print(tool_msg.content)

